#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import logging
import re
import json
import time
import os

import copy
from elasticsearch import Elasticsearch, NotFoundError
from elasticsearch_dsl import UpdateByQuery, Q, Search, Index
from elastic_transport import ConnectionTimeout
from rag import settings
from rag.settings import TAG_FLD, PAGERANK_FLD
from rag.utils import singleton, get_float
from api.utils.file_utils import get_project_base_directory
from rag.utils.doc_store_conn import DocStoreConnection, MatchExpr, OrderByExpr, MatchTextExpr, MatchDenseExpr, \
    FusionExpr
from rag.nlp import is_english, rag_tokenizer

ATTEMPT_TIME = 2

logger = logging.getLogger('ragflow.es_conn')


# 单例装饰器，确保该类在整个应用中只有一个实例存在
@singleton
# 定义 ESConnection 类，继承自 DocStoreConnection
class ESConnection(DocStoreConnection):
    # 初始化方法，在创建实例时自动调用
    def __init__(self):
        # 初始化 info 属性为一个空字典，用于存储 Elasticsearch 的信息
        self.info = {}
        # 记录日志，显示使用的是哪个 Elasticsearch 主机作为文档引擎
        logger.info(f"Use Elasticsearch {settings.ES['hosts']} as the doc engine.")
        # 循环尝试连接 Elasticsearch 多次（次数由 ATTEMPT_TIME 决定）
        for _ in range(ATTEMPT_TIME):
            # 尝试建立 Elasticsearch 连接
            try:
                # 创建 Elasticsearch 实例，传入主机地址、认证信息、证书验证设置和超时时间
                self.es = Elasticsearch(
                    settings.ES["hosts"].split(","),
                    basic_auth=(settings.ES["username"], settings.ES[
                        "password"]) if "username" in settings.ES and "password" in settings.ES else None,
                    verify_certs=False,
                    timeout=600
                )
                # 如果成功创建了 Elasticsearch 实例，则获取其信息并跳出循环
                if self.es:
                    self.info = self.es.info()
                    break
            # 捕获任何异常，并记录警告日志，等待 5 秒后继续下一次尝试
            except Exception as e:
                logger.warning(f"{str(e)}. Waiting Elasticsearch {settings.ES['hosts']} to be healthy.")
                time.sleep(5)
        # 如果最终无法 ping 通 Elasticsearch，则记录错误日志并抛出异常
        if not self.es.ping():
            msg = f"Elasticsearch {settings.ES['hosts']} is unhealthy in 120s."
            logger.error(msg)
            raise Exception(msg)
        # 获取 Elasticsearch 的版本号，并提取主版本号
        v = self.info.get("version", {"number": "8.11.3"})
        v = v["number"].split(".")[0]
        # 如果主版本号小于 8，则记录错误日志并抛出异常
        if int(v) < 8:
            msg = f"Elasticsearch version must be greater than or equal to 8, current version: {v}"
            logger.error(msg)
            raise Exception(msg)
        # 构建映射文件的路径
        fp_mapping = os.path.join(get_project_base_directory(), "conf", "mapping.json")
        # 如果映射文件不存在，则记录错误日志并抛出异常
        if not os.path.exists(fp_mapping):
            msg = f"Elasticsearch mapping file not found at {fp_mapping}"
            logger.error(msg)
            raise Exception(msg)
        # 加载映射文件内容到 self.mapping 中
        self.mapping = json.load(open(fp_mapping, "r"))
        # 记录日志，表示 Elasticsearch 已经健康运行
        logger.info(f"Elasticsearch {settings.ES['hosts']} is healthy.")

    """
    Database operations
    """

    # 返回数据库类型为 elasticsearch
    def dbType(self) -> str:
        return "elasticsearch"

    # 返回集群的健康状态，并标记类型为 elasticsearch
    def health(self) -> dict:
        health_dict = dict(self.es.cluster.health())
        health_dict["type"] = "elasticsearch"
        return health_dict

    """
    Table operations
    """

    # 创建索引，如果索引已存在则直接返回 True
    def createIdx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
        if self.indexExist(indexName, knowledgebaseId):
            return True
        # 导入 IndicesClient 并尝试创建索引
        try:
            from elasticsearch.client import IndicesClient
            return IndicesClient(self.es).create(index=indexName,
                                                 settings=self.mapping["settings"],
                                                 mappings=self.mapping["mappings"])
        # 捕获异常并记录日志
        except Exception:
            logger.exception("ESConnection.createIndex error %s" % (indexName))

    # 删除索引，但如果 knowledgebaseId 非空则不删除
    def deleteIdx(self, indexName: str, knowledgebaseId: str):
        if len(knowledgebaseId) > 0:
            # 索引需要在任何知识库删除后仍然存在，因为这个租户下的所有知识库都在同一个索引中
            # The index need to be alive after any kb deletion since all kb under this tenant are in one index.
            return
        # 尝试删除索引，允许无索引的情况
        try:
            self.es.indices.delete(index=indexName, allow_no_indices=True)
        # 忽略 NotFoundError 异常
        except NotFoundError:
            pass
        # 捕获其他异常并记录日志
        except Exception:
            logger.exception("ESConnection.deleteIdx error %s" % (indexName))

    # 检查索引是否存在，最多尝试 ATTEMPT_TIME 次
    def indexExist(self, indexName: str, knowledgebaseId: str = None) -> bool:
        s = Index(indexName, self.es)
        # 最多尝试 ATTEMPT_TIME 次检查索引是否存在
        for i in range(ATTEMPT_TIME):
            try:
                return s.exists()
            except Exception as e:
                # 记录异常日志
                logger.exception("ESConnection.indexExist got exception")
                # 如果是 Timeout 或 Conflict 异常，则继续尝试
                if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0:
                    continue
                # 否则中断循环并返回 False
                break
        return False

    """
    CRUD operations
    """

    # 定义 search 方法，用于执行 Elasticsearch 的搜索操作，并接收多个参数用于构建查询条件和返回格式
    def search(
            self, selectFields: list[str],
            highlightFields: list[str],
            condition: dict,
            matchExprs: list[MatchExpr],
            orderBy: OrderByExpr,
            offset: int,
            limit: int,
            indexNames: str | list[str],
            knowledgebaseIds: list[str],
            aggFields: list[str] = [],
            rank_feature: dict | None = None
    ):
        """
        Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html
        """
        # 如果 indexNames 是字符串类型，则按逗号分割为列表
        if isinstance(indexNames, str):
            indexNames = indexNames.split(",")
        # 确保 indexNames 是非空列表
        assert isinstance(indexNames, list) and len(indexNames) > 0
        # 确保 condition 中不包含 "_id" 字段
        assert "_id" not in condition

        # 初始化一个空的布尔查询对象，后续将添加过滤和匹配条件
        bqry = Q("bool", must=[])
        # 将知识库 ID 添加到查询条件中
        condition["kb_id"] = knowledgebaseIds
        # 遍历 condition 字典，构建过滤条件
        for k, v in condition.items():
            # 处理 available_int 字段的特殊逻辑：等于 0 或大于 0
            if k == "available_int":
                if v == 0:
                    bqry.filter.append(Q("range", available_int={"lt": 1}))
                else:
                    bqry.filter.append(
                        Q("bool", must_not=Q("range", available_int={"lt": 1})))
                continue
            # 如果值为空，则跳过该字段的条件
            if not v:
                continue
            # 如果值是列表类型，则使用 terms 查询
            if isinstance(v, list):
                bqry.filter.append(Q("terms", **{k: v}))
            # 如果值是字符串或整数类型，则使用 term 查询
            elif isinstance(v, str) or isinstance(v, int):
                bqry.filter.append(Q("term", **{k: v}))
            # 否则抛出异常，表示不支持的值类型
            else:
                raise Exception(
                    f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")

        # 创建一个新的 Search 对象
        s = Search()
        # 默认向量相似度权重设为 0.5
        vector_similarity_weight = 0.5
        # 第一次遍历 matchExprs，查找 FusionExpr 类型并设置权重
        for m in matchExprs:
            if isinstance(m, FusionExpr) and m.method == "weighted_sum" and "weights" in m.fusion_params:
                # 断言确保 matchExprs 包含三个特定类型的表达式
                assert len(matchExprs) == 3 and isinstance(matchExprs[0], MatchTextExpr) and isinstance(matchExprs[1],
                                                                                                        MatchDenseExpr) and isinstance(
                    matchExprs[2], FusionExpr)
                weights = m.fusion_params["weights"]
                # 提取向量相似度权重
                vector_similarity_weight = get_float(weights.split(",")[1])
        # 第二次遍历 matchExprs，分别处理文本匹配和向量匹配
        for m in matchExprs:
            if isinstance(m, MatchTextExpr):
                # 获取 minimum_should_match 参数，默认为 0.0
                minimum_should_match = m.extra_options.get("minimum_should_match", 0.0)
                if isinstance(minimum_should_match, float):
                    minimum_should_match = str(int(minimum_should_match * 100)) + "%"
                # 添加文本匹配查询条件
                bqry.must.append(Q("query_string", fields=m.fields,
                                   type="best_fields", query=m.matching_text,
                                   minimum_should_match=minimum_should_match,
                                   boost=1))
                # 设置整体查询的 boost 权重
                bqry.boost = 1.0 - vector_similarity_weight

            elif isinstance(m, MatchDenseExpr):
                assert (bqry is not None)
                # 获取相似度阈值，默认为 0.0
                similarity = 0.0
                if "similarity" in m.extra_options:
                    similarity = m.extra_options["similarity"]
                # 使用 knn 方法进行向量相似性搜索
                s = s.knn(m.vector_column_name,
                          m.topn,
                          m.topn * 2,
                          query_vector=list(m.embedding_data),
                          filter=bqry.to_dict(),
                          similarity=similarity,
                          )

        # 如果存在 rank_feature，则将其作为 should 条件加入查询
        if bqry and rank_feature:
            for fld, sc in rank_feature.items():
                if fld != PAGERANK_FLD:
                    # 构建嵌套字段名称
                    fld = f"{TAG_FLD}.{fld}"
                # 添加 rank_feature 查询条件
                bqry.should.append(Q("rank_feature", field=fld, linear={}, boost=sc))

        # 如果有查询条件，则将其应用到搜索对象上
        if bqry:
            s = s.query(bqry)
        # 为指定字段添加高亮显示
        for field in highlightFields:
            s = s.highlight(field)

        # 处理排序逻辑
        if orderBy:
            orders = list()
            for field, order in orderBy.fields:
                # 将排序顺序转换为 asc 或 desc
                order = "asc" if order == 0 else "desc"
                # 根据字段类型设置不同的排序选项
                if field in ["page_num_int", "top_int"]:
                    order_info = {"order": order, "unmapped_type": "float",
                                  "mode": "avg", "numeric_type": "double"}
                elif field.endswith("_int") or field.endswith("_flt"):
                    order_info = {"order": order, "unmapped_type": "float"}
                else:
                    order_info = {"order": order, "unmapped_type": "text"}
                orders.append({field: order_info})
            # 应用排序规则
            s = s.sort(*orders)

        # 为每个聚合字段添加 terms 聚合
        for fld in aggFields:
            s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)

        # 设置分页参数
        if limit > 0:
            s = s[offset:offset + limit]
        # 将查询对象转换为字典形式
        q = s.to_dict()
        # 记录生成的查询语句（调试级别日志）
        logger.debug(f"ESConnection.search {str(indexNames)} query: " + json.dumps(q))

        # 最多重试 ATTEMPT_TIME 次执行搜索
        for i in range(ATTEMPT_TIME):
            try:
                #print(json.dumps(q, ensure_ascii=False))
                # 执行 Elasticsearch 搜索请求
                res = self.es.search(index=indexNames,
                                     body=q,
                                     timeout="600s",
                                     # search_type="dfs_query_then_fetch",
                                     track_total_hits=True,
                                     _source=True)
                # 如果搜索超时，抛出异常
                if str(res.get("timed_out", "")).lower() == "true":
                    raise Exception("Es Timeout.")
                # 记录搜索结果（调试级别日志）
                logger.debug(f"ESConnection.search {str(indexNames)} res: " + str(res))
                # 返回搜索结果
                return res
            except Exception as e:
                # 记录异常信息和查询内容
                logger.exception(f"ESConnection.search {str(indexNames)} query: " + str(q))
                # 如果是超时异常，继续尝试
                if str(e).find("Timeout") > 0:
                    continue
                # 其他异常直接抛出
                raise e
        # 如果三次都超时，记录错误并抛出异常
        logger.error("ESConnection.search timeout for 3 times!")
        raise Exception("ESConnection.search timeout.")


    # 定义 get 方法，用于根据 chunkId 从 Elasticsearch 中获取一个文档
    def get(self, chunkId: str, indexName: str, knowledgebaseIds: list[str]) -> dict | None:
        # 最多重试 ATTEMPT_TIME 次执行获取操作
        for i in range(ATTEMPT_TIME):
            try:
                # 执行 Elasticsearch 的 get 请求以获取指定 ID 的文档
                res = self.es.get(index=(indexName),
                                  id=chunkId, source=True, )
                # 如果搜索超时，抛出异常
                if str(res.get("timed_out", "")).lower() == "true":
                    raise Exception("Es Timeout.")
                # 提取文档内容，并添加 chunkId 字段后返回
                chunk = res["_source"]
                chunk["id"] = chunkId
                return chunk
            # 如果文档不存在，则返回 None
            except NotFoundError:
                return None
            # 捕获其他异常并记录日志
            except Exception as e:
                logger.exception(f"ESConnection.get({chunkId}) got exception")
                # 如果是超时异常，则继续尝试
                if str(e).find("Timeout") > 0:
                    continue
                # 其他异常直接抛出
                raise e
        # 如果三次都超时，记录错误并抛出异常
        logger.error("ESConnection.get timeout for 3 times!")
        raise Exception("ESConnection.get timeout.")

    # 定义 insert 方法，用于将多个文档批量插入到 Elasticsearch 索引中
    def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str = None) -> list[str]:
        # Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
        # 初始化空的操作列表
        operations = []
        # 遍历 documents 列表，构建 bulk 插入所需的命令结构
        for d in documents:
            # 确保每个文档不包含 _id 字段，但必须包含 id 字段
            assert "_id" not in d
            assert "id" in d
            # 深拷贝当前文档，并为其添加 kb_id 字段
            d_copy = copy.deepcopy(d)
            d_copy["kb_id"] = knowledgebaseId
            # 移除 id 字段，并将其作为 meta_id 使用
            meta_id = d_copy.pop("id", "")
            # 添加 index 操作指令和对应的文档数据
            operations.append(
                {"index": {"_index": indexName, "_id": meta_id}})
            operations.append(d_copy)

        # 初始化结果列表
        res = []
        # 最多重试 ATTEMPT_TIME 次执行 bulk 插入操作
        for _ in range(ATTEMPT_TIME):
            try:
                # 初始化结果列表
                res = []
                # 执行 Elasticsearch 的 bulk 插入操作
                r = self.es.bulk(index=(indexName), operations=operations,
                                 refresh=False, timeout="60s")
                # 如果响应中没有错误，则返回空结果（表示成功）
                if re.search(r"False", str(r["errors"]), re.IGNORECASE):
                    return res

                # 遍历响应中的每条记录，提取可能的错误信息
                for item in r["items"]:
                    for action in ["create", "delete", "index", "update"]:
                        if action in item and "error" in item[action]:
                            res.append(str(item[action]["_id"]) + ":" + str(item[action]["error"]))
                # 返回错误信息列表
                return res
            # 捕获异常并记录警告日志
            except Exception as e:
                res.append(str(e))
                logger.warning("ESConnection.insert got exception: " + str(e))
                res = []
                # 如果是超时或类似超时的异常，则等待 3 秒后继续尝试
                if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE):
                    res.append(str(e))
                    time.sleep(3)
                    continue
        # 返回最终的结果列表
        return res

    # 定义 update 方法，用于根据条件更新 Elasticsearch 中的文档数据
    def update(self, condition: dict, newValue: dict, indexName: str, knowledgebaseId: str) -> bool:
        # 深拷贝新值数据，并移除 id 字段
        doc = copy.deepcopy(newValue)
        doc.pop("id", None)
        # 将知识库 ID 添加到查询条件中
        condition["kb_id"] = knowledgebaseId
        # 如果条件中包含 id 字段且为字符串类型，则更新特定的单个文档
        if "id" in condition and isinstance(condition["id"], str):
            # 获取 chunkId 以进行更新
            # update specific single document
            chunkId = condition["id"]
            # 最多重试 ATTEMPT_TIME 次执行更新操作
            for i in range(ATTEMPT_TIME):
                # 遍历文档字段，删除名称后缀为 "feas" 的字段
                for k in doc.keys():
                    if "feas" != k.split("_")[-1]:
                        continue
                    try:
                        self.es.update(index=indexName, id=chunkId, script=f"ctx._source.remove(\"{k}\");")
                    except Exception:
                        logger.exception(f"ESConnection.update(index={indexName}, id={chunkId}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
                try:
                    # 更新文档内容
                    self.es.update(index=indexName, id=chunkId, doc=doc)
                    return True
                except Exception as e:
                    # 记录异常信息
                    logger.exception(
                        f"ESConnection.update(index={indexName}, id={chunkId}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
                    # 如果是超时或连接异常，则继续尝试
                    if re.search(r"(timeout|connection)", str(e).lower()):
                        continue
                    break
            # 更新失败，返回 False
            return False

        # update unspecific maybe-multiple documents
        # 构建布尔查询对象，用于非特定文档更新（可能多条记录）
        bqry = Q("bool")
        # 遍历 condition 字典，构建过滤条件
        for k, v in condition.items():
            # 忽略无效的键值对
            if not isinstance(k, str) or not v:
                continue
            # 处理 exists 条件：判断字段是否存在
            if k == "exists":
                bqry.filter.append(Q("exists", field=v))
                continue
            # 如果值是列表类型，则使用 terms 查询
            if isinstance(v, list):
                bqry.filter.append(Q("terms", **{k: v}))
            # 如果值是字符串或整数类型，则使用 term 查询
            elif isinstance(v, str) or isinstance(v, int):
                bqry.filter.append(Q("term", **{k: v}))
            # 否则抛出异常，表示不支持的值类型
            else:
                raise Exception(
                    f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
        # 初始化 scripts 和 params 列表，用于存储脚本更新逻辑和参数
        scripts = []
        params = {}
        # 遍历 newValue 字典，构建更新脚本
        for k, v in newValue.items():
            # 处理 remove 操作：从文档中移除指定字段或元素
            if k == "remove":
                if isinstance(v, str):
                    scripts.append(f"ctx._source.remove('{v}');")
                if isinstance(v, dict):
                    for kk, vv in v.items():
                        scripts.append(f"int i=ctx._source.{kk}.indexOf(params.p_{kk});ctx._source.{kk}.remove(i);")
                        params[f"p_{kk}"] = vv
                continue
            # 处理 add 操作：向指定字段添加新值
            if k == "add":
                if isinstance(v, dict):
                    for kk, vv in v.items():
                        scripts.append(f"ctx._source.{kk}.add(params.pp_{kk});")
                        params[f"pp_{kk}"] = vv.strip()
                continue
            # 忽略无效的键值对（除了 available_int）
            if (not isinstance(k, str) or not v) and k != "available_int":
                continue
            # 处理字符串类型的字段更新
            if isinstance(v, str):
                v = re.sub(r"(['\n\r]|\\.)", " ", v)
                params[f"pp_{k}"] = v
                scripts.append(f"ctx._source.{k}=params.pp_{k};")
            # 处理数值类型的字段更新
            elif isinstance(v, int) or isinstance(v, float):
                scripts.append(f"ctx._source.{k}={v};")
            # 处理列表类型的字段更新
            elif isinstance(v, list):
                scripts.append(f"ctx._source.{k}=params.pp_{k};")
                params[f"pp_{k}"] = json.dumps(v, ensure_ascii=False)
            # 否则抛出异常，表示不支持的值类型
            else:
                raise Exception(
                    f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
        # 创建 UpdateByQuery 对象并设置相关参数
        ubq = UpdateByQuery(
            index=indexName).using(
            self.es).query(bqry)
        ubq = ubq.script(source="".join(scripts), params=params)
        ubq = ubq.params(refresh=True)
        ubq = ubq.params(slices=5)
        ubq = ubq.params(conflicts="proceed")

        # 最多重试 ATTEMPT_TIME 次执行批量更新操作
        for _ in range(ATTEMPT_TIME):
            try:
                # 执行批量更新并返回结果
                _ = ubq.execute()
                return True
            except Exception as e:
                # 记录错误日志和执行脚本
                logger.error("ESConnection.update got exception: " + str(e) + "\n".join(scripts))
                # 如果是超时、连接或冲突异常，则继续尝试
                if re.search(r"(timeout|connection|conflict)", str(e).lower()):
                    continue
                break
        # 更新失败，返回 False
        return False

    # 定义 delete 方法，用于根据条件删除 Elasticsearch 中的文档
    def delete(self, condition: dict, indexName: str, knowledgebaseId: str) -> int:
        # 初始化查询对象
        qry = None
        # 确保 condition 不包含 "_id" 字段
        assert "_id" not in condition
        # 将知识库 ID 添加到查询条件中
        condition["kb_id"] = knowledgebaseId
        # 如果 condition 包含 id 字段，则构造基于 ids 的查询
        if "id" in condition:
            chunk_ids = condition["id"]
            # 确保 chunk_ids 是列表类型
            if not isinstance(chunk_ids, list):
                chunk_ids = [chunk_ids]
            # 如果 chunk_ids 为空，则删除所有文档
            if not chunk_ids:  # when chunk_ids is empty, delete all
                qry = Q("match_all")
            else:
                # 根据 id 列表构建查询条件
                qry = Q("ids", values=chunk_ids)
        else:
            # 构建更复杂的查询条件
            qry = Q("bool")
            for k, v in condition.items():
                # 处理 exists 条件：判断字段是否存在
                if k == "exists":
                    qry.filter.append(Q("exists", field=v))

                # 处理 must_not 条件：例如判断字段不存在
                elif k == "must_not":
                    if isinstance(v, dict):
                        for kk, vv in v.items():
                            if kk == "exists":
                                qry.must_not.append(Q("exists", field=vv))

                # 处理列表类型的匹配条件
                elif isinstance(v, list):
                    qry.must.append(Q("terms", **{k: v}))
                # 处理字符串或整数类型的匹配条件
                elif isinstance(v, str) or isinstance(v, int):
                    qry.must.append(Q("term", **{k: v}))
                # 否则抛出异常，表示不支持的值类型
                else:
                    raise Exception("Condition value must be int, str or list.")
        # 记录生成的查询语句（调试级别日志）
        logger.debug("ESConnection.delete query: " + json.dumps(qry.to_dict()))
        # 最多重试 ATTEMPT_TIME 次执行删除操作
        for _ in range(ATTEMPT_TIME):
            try:
                #print(Search().query(qry).to_dict(), flush=True)
                # 执行 delete_by_query 请求删除符合条件的文档
                res = self.es.delete_by_query(
                    index=indexName,
                    body=Search().query(qry).to_dict(),
                    refresh=True)
                # 返回已删除的文档数量
                return res["deleted"]
            except Exception as e:
                # 记录警告日志
                logger.warning("ESConnection.delete got exception: " + str(e))
                # 如果是超时或连接异常，则等待 3 秒后继续尝试
                if re.search(r"(timeout|connection)", str(e).lower()):
                    time.sleep(3)
                    continue
                # 如果是未找到异常，则返回 0 表示没有删除任何文档
                if re.search(r"(not_found)", str(e), re.IGNORECASE):
                    return 0
        # 删除失败，返回 0
        return 0

    """
    Helper functions for search result
    """

    # 解析 Elasticsearch 返回结果中的总命中数
    def getTotal(self, res):
        # 判断 total 字段是否为字典类型（Elasticsearch 7+ 的新格式）
        if isinstance(res["hits"]["total"], type({})):
            # 如果是字典，返回 "value" 字段的值
            return res["hits"]["total"]["value"]
        # 否则直接返回 total 值（旧版本格式）
        return res["hits"]["total"]

    # 从 Elasticsearch 查询结果中提取所有文档的 ID
    def getChunkIds(self, res):
        # 遍历 hits 中的所有文档，提取 _id 字段组成列表返回
        return [d["_id"] for d in res["hits"]["hits"]]

    # 私有方法：从 Elasticsearch 结果中提取源数据并附加文档 ID 和得分
    def __getSource(self, res):
        # 初始化结果列表
        rr = []
        # 遍历每个命中项，将 _source 数据加上 _id 和 _score 组合成新对象并加入结果列表
        for d in res["hits"]["hits"]:
            d["_source"]["id"] = d["_id"]
            d["_source"]["_score"] = d["_score"]
            rr.append(d["_source"])
        # 返回处理后的源数据列表
        return rr

    # 根据指定字段列表从查询结果中提取字段内容，返回以 chunkId 为键的字典
    def getFields(self, res, fields: list[str]) -> dict[str, dict]:
        # 初始化结果字典
        res_fields = {}
        # 如果字段列表为空，直接返回空字典
        if not fields:
            return {}
        # 遍历处理后的源数据
        for d in self.__getSource(res):
            # 提取当前文档中包含的指定字段，并过滤掉值为 None 的字段
            m = {n: d.get(n) for n in fields if d.get(n) is not None}
            # 遍历提取出的字段，进行类型判断和转换
            for n, v in m.items():
                if isinstance(v, list):
                    # 如果值是列表类型，保留原样
                    m[n] = v
                    continue
                if not isinstance(v, str):
                    # 如果不是字符串类型，转为字符串
                    m[n] = str(m[n])
                # if n.find("tks") > 0:
                #     m[n] = rmSpace(m[n])

            # 如果提取到字段内容非空，则将其按文档 ID 存入结果字典
            if m:
                res_fields[d["id"]] = m
        # 返回最终提取的字段字典
        return res_fields

    # 从 Elasticsearch 响应中提取高亮内容，并根据关键词进行高亮处理
    def getHighlight(self, res, keywords: list[str], fieldnm: str):
        # 初始化结果字典
        ans = {}
        # 遍历每个命中项，提取高亮字段内容
        for d in res["hits"]["hits"]:
            hlts = d.get("highlight")
            if not hlts:
                continue
            # 将高亮内容拼接成字符串
            txt = "...".join([a for a in list(hlts.items())[0][1]])
            # 如果不是英文文本，则直接存入结果并跳过后续处理
            if not is_english(txt.split()):
                ans[d["_id"]] = txt
                continue

            # 否则获取源字段内容并进行正则处理
            txt = d["_source"][fieldnm]
            txt = re.sub(r"[\r\n]", " ", txt, flags=re.IGNORECASE | re.MULTILINE)
            txts = []
            # 按句子拆分，并对每个关键词进行匹配替换（添加 <em> 标签）
            for t in re.split(r"[.?!;\n]", txt):
                for w in keywords:
                    t = re.sub(r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w), r"\1<em>\2</em>\3", t,
                               flags=re.IGNORECASE | re.MULTILINE)
                # 如果没有找到匹配的高亮内容，则跳过该句子
                if not re.search(r"<em>[^<>]+</em>", t, flags=re.IGNORECASE | re.MULTILINE):
                    continue
                txts.append(t)
            # 将处理后的句子拼接后存入结果字典，若无有效高亮则使用原始高亮内容
            ans[d["_id"]] = "...".join(txts) if txts else "...".join([a for a in list(hlts.items())[0][1]])

        # 返回最终的高亮结果
        return ans

    # 提取指定字段的聚合统计结果（如 buckets 中的 key 和 doc_count）
    def getAggregation(self, res, fieldnm: str):
        # 构造聚合字段名称
        agg_field = "aggs_" + fieldnm
        # 如果聚合结果不存在，则返回空列表
        if "aggregations" not in res or agg_field not in res["aggregations"]:
            return list()
        # 获取 buckets 数据并提取 key 和 doc_count 组成元组列表返回
        bkts = res["aggregations"][agg_field]["buckets"]
        return [(b["key"], b["doc_count"]) for b in bkts]

    """
    SQL
    """

    # 定义 sql 方法用于执行类 SQL 查询，并做了一些语法转换和重试机制
    def sql(self, sql: str, fetch_size: int, format: str):
        # 记录原始 SQL 到 debug 日志
        logger.debug(f"ESConnection.sql get sql: {sql}")
        # 替换多余的空白字符为单个空格
        sql = re.sub(r"[ `]+", " ", sql)
        # 移除 SQL 中的所有百分号
        sql = sql.replace("%", "")
        # 初始化替换列表
        replaces = []
        # 正则查找所有 like 或等于操作的字段和值
        for r in re.finditer(r" ([a-z_]+_l?tks)( like | ?= ?)'([^']+)'", sql):
            fld, v = r.group(1), r.group(3)
            # 构造新的 MATCH 查询语句
            match = " MATCH({}, '{}', 'operator=OR;minimum_should_match=30%') ".format(
                fld, rag_tokenizer.fine_grained_tokenize(rag_tokenizer.tokenize(v)))
            # 添加到替换列表
            replaces.append(
                ("{}{}'{}'".format(
                    r.group(1),
                    r.group(2),
                    r.group(3)),
                 match))

        # 执行所有 SQL 中的替换逻辑
        for p, r in replaces:
            sql = sql.replace(p, r, 1)
        # 记录转换后的 SQL 到 debug 日志
        logger.debug(f"ESConnection.sql to es: {sql}")

        # 最多重试 ATTEMPT_TIME 次执行 SQL 查询
        for i in range(ATTEMPT_TIME):
            try:
                # 执行 Elasticsearch SQL 查询
                res = self.es.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format,
                                        request_timeout="2s")
                # 返回查询结果
                return res
            except ConnectionTimeout:
                # 超时异常，记录日志并继续重试
                logger.exception("ESConnection.sql timeout")
                continue
            except Exception:
                # 其他异常，记录日志并返回 None
                logger.exception("ESConnection.sql got exception")
                return None
        # 三次都失败，记录错误并返回 None
        logger.error("ESConnection.sql timeout for 3 times!")
        return None
